import argparse
import json
import os
import sys

# 添加项目根目录到Python路径
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from configs.default_config import DEFAULT_CONFIG

def str2bool(v):
    """将字符串转换为布尔值"""
    if isinstance(v, bool):
        return v
    if v.lower() in ('yes', 'true', 't', 'y', '1'):
        return True
    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
        return False
    else:
        raise argparse.ArgumentTypeError('Boolean value expected.')

def parse_args():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description="PoXiao-LLM Training")
    
    # 基础配置
    parser.add_argument("--config_file", type=str, default=None, help="配置文件路径")
    parser.add_argument("--output_dir", type=str, default=DEFAULT_CONFIG["output_dir"], help="输出目录")
    parser.add_argument("--seed", type=int, default=DEFAULT_CONFIG["seed"], help="随机种子")
    
    # 训练配置
    parser.add_argument("--epochs", type=int, default=DEFAULT_CONFIG["epochs"], help="训练轮数")
    parser.add_argument("--batch_size", type=int, default=DEFAULT_CONFIG["batch_size"], help="批次大小")
    parser.add_argument("--learning_rate", type=float, default=DEFAULT_CONFIG["learning_rate"], help="学习率")
    parser.add_argument("--device", type=str, default=DEFAULT_CONFIG["device"], help="设备")
    parser.add_argument("--dtype", type=str, default=DEFAULT_CONFIG["dtype"], help="数据类型")
    parser.add_argument("--accumulation_steps", type=int, default=DEFAULT_CONFIG["accumulation_steps"], help="梯度累积步数")
    parser.add_argument("--grad_clip", type=float, default=DEFAULT_CONFIG["grad_clip"], help="梯度裁剪阈值")
    parser.add_argument("--warmup_iters", type=int, default=DEFAULT_CONFIG["warmup_iters"], help="预热步数")
    parser.add_argument("--log_interval", type=int, default=DEFAULT_CONFIG["log_interval"], help="日志间隔")
    parser.add_argument("--save_interval", type=int, default=DEFAULT_CONFIG["save_interval"], help="保存间隔")
    
    # 模型配置
    parser.add_argument("--hidden_size", type=int, default=DEFAULT_CONFIG["hidden_size"], help="隐藏层大小")
    parser.add_argument("--num_hidden_layers", type=int, default=DEFAULT_CONFIG["num_hidden_layers"], help="隐藏层数量")
    parser.add_argument("--max_seq_len", type=int, default=DEFAULT_CONFIG["max_seq_len"], help="最大序列长度")
    parser.add_argument("--use_moe", type=str2bool, default=DEFAULT_CONFIG["use_moe"], help="是否使用MoE")
    
    # 插件配置
    parser.add_argument("--model_plugin", type=str, default=DEFAULT_CONFIG["model_plugin"], help="模型插件")
    parser.add_argument("--dataset_plugin", type=str, default=DEFAULT_CONFIG["dataset_plugin"], help="数据集插件")
    parser.add_argument("--loss_plugin", type=str, default=DEFAULT_CONFIG["loss_plugin"], help="损失函数插件")
    parser.add_argument("--optimizer_plugin", type=str, default=DEFAULT_CONFIG["optimizer_plugin"], help="优化器插件")
    parser.add_argument("--trainer_plugin", type=str, default=DEFAULT_CONFIG["trainer_plugin"], help="训练器插件")
    parser.add_argument("--scheduler_plugin", type=str, default=DEFAULT_CONFIG["scheduler_plugin"], help="学习率调度器插件")
    parser.add_argument("--checkpointer_plugin", type=str, default=DEFAULT_CONFIG["checkpointer_plugin"], help="保存点插件")
    
    # 早停法配置
    parser.add_argument("--early_stopping_patience", type=int, default=DEFAULT_CONFIG["early_stopping_patience"], help="早停法耐心次数")
    parser.add_argument("--early_stopping_min_delta", type=float, default=DEFAULT_CONFIG["early_stopping_min_delta"], help="早停法最小改善阈值")
    
    # 数据配置
    parser.add_argument("--data_path", type=str, default=DEFAULT_CONFIG["data_path"], help="数据路径")
    parser.add_argument("--tokenizer_path", type=str, default=DEFAULT_CONFIG["tokenizer_path"], help="分词器路径")
    
    # 分布式训练配置
    parser.add_argument("--use_accelerate", action='store_true', help="启用Accelerate")
    parser.add_argument("--use_torchrun", action='store_true', help="启用TorchRun")
    parser.add_argument("--num_workers", type=int, default=DEFAULT_CONFIG["num_workers"], help="数据加载器工作线程数")
    parser.add_argument("--self_schedule", type=str2bool, default=DEFAULT_CONFIG["self_schedule"], help="启用自定义学习率调度")
    
    # 日志配置
    parser.add_argument("--use_wandb", action='store_true', help="启用WandB日志")
    parser.add_argument("--use_tensorboard", action='store_true', help="启用TensorBoard日志")
    parser.add_argument("--tensorboard_log_dir", type=str, default=DEFAULT_CONFIG["tensorboard_log_dir"], help="TensorBoard日志目录")
    parser.add_argument("--wandb_project", type=str, default=DEFAULT_CONFIG["wandb_project"], help="WandB项目名称")
    
    args = parser.parse_args()
    
    # 如果提供了配置文件，则加载配置文件
    config = DEFAULT_CONFIG.copy()
    if args.config_file:
        if os.path.exists(args.config_file):
            with open(args.config_file, 'r') as f:
                file_config = json.load(f)
                config.update(file_config)
        else:
            print(f"警告: 配置文件 {args.config_file} 不存在，将使用默认配置")
            
    # 用命令行参数覆盖配置
    for key, value in vars(args).items():
        if value is not None and value != DEFAULT_CONFIG.get(key):
            config[key] = value
            
    return config